
import cv2
import numpy

from Settings import BaseTestDataCubePath, ModelsBaseDIR
import gfpgan




class FacePictureEngine:
    def __init__(self):
        # 图像增强
        __model_path = f'{ModelsBaseDIR}/GFPGAN.pth'
        self. __face_engine = gfpgan.GFPGANer(
            model_path=__model_path,
            upscale=1
        )

    def enhancePicture(self, image_path) -> numpy.ndarray:
        _, enhance_face_image, _ = self.__face_engine.enhance(img=image_path)
        return enhance_face_image[0]

    def fixPicture(self, image_path):
        self. __face_engine.


if __name__ == '__main__':
    face_engine = FacePictureEngine()
    source_face_image_path = f'{BaseTestDataCubePath}/img.png'
    result = face_engine.enhancePicture(source_face_image_path)
    cv2.imwrite(f'{BaseTestDataCubePath}/4.jpg', result)
    # cv2.imshow('enhancePicture', result)



